"""
使用 Scikit-learn 构建垃圾邮件检测器
"""
import pandas as pd
import jieba                                                        # 中文分词
from sklearn.feature_extraction.text import TfidfVectorizer         # TF-IDF向量化
from sklearn.naive_bayes import MultinomialNB                       # 朴素贝叶斯分类器   
from sklearn.model_selection import train_test_split                # 数据集划分
from sklearn.metrics import classification_report, accuracy_score   # 评估指标
import re

# 中文停用词列表
CHINESE_STOP_WORDS = [
    '的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一', '一个', '上', '也', '很', '到', '说', '要', '去', '你',
    '会', '着', '没有', '看', '好', '自己', '这', '那', '来', '用', '他', '她', '它', '们', '个', '啊', '吧', '呢', '吗', '嗯',
    '哦', '哈', '嘿', '呀', '哟', '呵', '唉', '嗯嗯', '哼', '呜', '哇', '咦', '咯', '咚', '咔', '咳', '咪', '咭', '咻', '咿',
    '为', '而', '如', '因', '所以', '但是', '然而', '可是', '虽然', '尽管', '无论', '不管', '无', '没', '还', '又', '才', '便',
    '这样', '那样', '这里', '那里', '这些', '那些', '什么', '怎么', '为什么', '怎样', '多少', '几', '第', '每', '各', '另',
    '其', '其他', '别', '别的', '谁', '哪', '哪里', '哪些', '何', '何时', '何地', '何人', '如何', '多么', '何等', '多大'
]

def preprocess_chinese_text(text):
    """中文文本预处理"""
    # 移除数字、英文字母、特殊字符，保留中文和空格
    text = re.sub(r'[^\u4e00-\u9fff\s]', '', text)
    
    # 使用jieba进行中文分词
    words = jieba.lcut(text)
    
    # 过滤停用词和长度小于2的词
    filtered_words = [word for word in words if word not in CHINESE_STOP_WORDS and len(word) >= 2]
    
    return ' '.join(filtered_words)

def main():
    # 创建中文垃圾邮件示例数据
    data = {
        'text': [
            # 垃圾邮件样本
            "恭喜您！您已中奖100万元！立即拨打电话领取大奖！千万不要错过！",
            "免费获得减肥药！一周瘦20斤！效果保证！不瘦退款！立即订购！",
            "紧急通知！您的银行账户即将被冻结！请立即点击链接验证身份！",
            "投资理财！月收益30%！无风险！保本保息！立即加入！",
            "免费送苹果手机！限时优惠！点击领取！机会难得！",
            "代开发票！增值税普票专票！价格优惠！质量保证！",
            "网络兼职！在家赚钱！日入300元！无需经验！",
            "股票内幕消息！明日涨停板！千载难逢机会！立即购买！",
            "办理贷款！无抵押！当天放款！利息低！额度高！",
            "减肥神器！一个月瘦30斤！明星同款！效果惊人！",
            "免费试用化妆品！美白祛斑！立竿见影！包邮到家！",
            "彩票预测！百分百中奖！专业团队！成功率极高！",
            
            # 正常邮件样本
            "明天下午3点开会，请大家准时参加。会议室在A栋302。",
            "您好，您的快递已经到达，请及时取件。快递单号：SF123456789。",
            "生日快乐！祝您身体健康，工作顺利，家庭幸福！",
            "会议纪要已发送，请查收附件并于明日前反馈意见。",
            "感谢您参加今天的培训，课件资料已上传到公司网盘。",
            "您的机票预订成功，航班号CA1234，明天早上8点起飞。",
            "项目进度汇报：目前已完成60%，预计下周完成全部工作。",
            "年终总结报告已提交，请领导审核批准。",
            "客户反馈很好，对我们的服务非常满意。",
            "明天的聚餐改到晚上7点，地点还是老地方。",
            "报销单据已整理完毕，请财务部门审核。",
            "产品发布会定于下周五举行，请提前准备相关材料。"
        ],
        'label': ['spam'] * 12 + ['ham'] * 12  # 前12个是垃圾邮件，后12个是正常邮件
    }
    
    df = pd.DataFrame(data)
    
    print("原始数据集信息：")
    print(f"总邮件数量: {len(df)}")
    print(f"垃圾邮件数量: {sum(df['label'] == 'spam')}")
    print(f"正常邮件数量: {sum(df['label'] == 'ham')}")
    print("\n" + "="*50 + "\n")
    
    # 预处理文本
    print("正在进行中文文本预处理...")
    df['processed_text'] = df['text'].apply(preprocess_chinese_text)
    
    # 显示预处理结果示例
    print("文本预处理示例：")
    for i in range(2):
        print(f"原文: {df.iloc[i]['text']}")
        print(f"预处理后: {df.iloc[i]['processed_text']}")
        print()
    
    # 创建TF-IDF特征（不使用英文停用词）
    vectorizer = TfidfVectorizer(
        max_features=1000,      # 最大特征数
        ngram_range=(1, 2),     # 使用1-gram和2-gram
        min_df=1,               # 最小文档频率
        max_df=0.8              # 最大文档频率
    )
    
    X = vectorizer.fit_transform(df['processed_text'])
    y = df['label'].map({'ham': 0, 'spam': 1})
    
    print(f"特征矩阵形状: {X.shape}")
    print(f"标签分布: spam={sum(y)}, ham={len(y)-sum(y)}")
    print("\n" + "="*50 + "\n")
    
    # 分割数据
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.3, random_state=42, stratify=y
    )
    
    print(f"训练集大小: {X_train.shape[0]}")
    print(f"测试集大小: {X_test.shape[0]}")
    
    # 训练朴素贝叶斯模型
    print("\n正在训练朴素贝叶斯模型...")
    model = MultinomialNB(alpha=1.0)
    model.fit(X_train, y_train)
    
    # 预测和评估
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    
    print(f"\n模型评估结果:")
    print(f"准确率: {accuracy:.4f}")
    print("\n详细分类报告:")
    print(classification_report(y_test, y_pred, target_names=['正常邮件', '垃圾邮件']))
    
    # 测试新的中文邮件
    test_emails = [
        "恭喜您中奖！请立即领取100万现金大奖！",
        "明天的会议延期到下周，具体时间另行通知。",
        "免费送iPhone！点击链接立即领取！",
        "您的订单已发货，预计明天送达。",
        "股票内幕消息！保证赚钱！立即购买！"
    ]
    
    print("\n" + "="*50)
    print("测试新邮件:")
    print("="*50)
    
    for i, email in enumerate(test_emails, 1):
        # 预处理测试邮件
        processed_test = preprocess_chinese_text(email)
        test_vector = vectorizer.transform([processed_test])
        
        # 预测
        prediction = model.predict(test_vector)[0]
        probability = model.predict_proba(test_vector)[0]
        
        print(f"\n邮件 {i}: {email}")
        print(f"预处理后: {processed_test}")
        print(f"预测结果: {'垃圾邮件' if prediction == 1 else '正常邮件'}")
        print(f"垃圾邮件概率: {probability[1]:.4f}")
        print(f"正常邮件概率: {probability[0]:.4f}")
        
        # 显示置信度
        confidence = max(probability)
        confidence_level = "高" if confidence > 0.8 else "中" if confidence > 0.6 else "低"
        print(f"置信度: {confidence:.4f} ({confidence_level})")
    
    # 显示最重要的特征词
    print("\n" + "="*50)
    print("最重要的垃圾邮件特征词:")
    print("="*50)
    
    feature_names = vectorizer.get_feature_names_out()
    spam_log_probs = model.feature_log_prob_[1]  # 垃圾邮件类别的对数概率
    
    # 获取概率最高的前20个特征
    top_spam_indices = spam_log_probs.argsort()[-20:][::-1]
    
    print("垃圾邮件关键词 (按重要性排序):")
    for idx in top_spam_indices:
        feature = feature_names[idx]
        prob = spam_log_probs[idx]
        print(f"  {feature}: {prob:.4f}")

if __name__ == "__main__":
    # 在运行前，需要安装jieba库
    try:
        import jieba
    except ImportError:
        print("请先安装jieba库: pip install jieba")
        exit(1)
    
    main()